"""
Fortran AST builder for CodeMCP framework.

This module builds abstract syntax trees from Fortran source code
using fparser2 for robust parsing of modern and legacy Fortran.
"""

from typing import Dict, List, Any, Optional, Union
from pathlib import Path
import logging

try:
    from fparser2.parser.fortran2003 import FortranReader, FortranSyntaxError
    from fparser2.parser import FortranFileReader, FortranStringReader
    from fparser2.parser.two.parser import FortranReader as F2003Reader
    from fparser2.common.readfortran import Comment
    from fparser2.parser.two import Fortran2003
    FPARSER2_AVAILABLE = True
except ImportError:
    FPARSER2_AVAILABLE = False
    logging.warning("fparser2 not available. Fortran parsing will be limited.")


class FortranASTNode:
    """Base class for Fortran AST nodes."""
    
    def __init__(self, node_type: str, children: List = None, **attributes):
        self.node_type = node_type
        self.children = children or []
        self.attributes = attributes
        self.line_number = attributes.get('line_number', 0)
        self.source_text = attributes.get('source_text', '')
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert node to dictionary representation."""
        return {
            'type': self.node_type,
            'children': [child.to_dict() if hasattr(child, 'to_dict') else str(child) 
                        for child in self.children],
            'attributes': self.attributes,
            'line_number': self.line_number,
            'source_text': self.source_text
        }


class FortranASTBuilder:
    """
    Builds Fortran AST from source code using fparser2.
    
    Provides structured representation of Fortran programs including:
    - Program units (programs, modules, procedures)
    - Type definitions and derived types
    - Interface blocks and generic procedures
    - Variable declarations and specifications
    - Executable constructs and statements
    """
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        if not FPARSER2_AVAILABLE:
            self.logger.error("fparser2 is required for Fortran AST building")
    
    def build_from_file(self, file_path: str) -> FortranASTNode:
        """
        Build AST from Fortran source file.
        
        Args:
            file_path: Path to Fortran source file
            
        Returns:
            Root AST node
            
        Raises:
            FileNotFoundError: If file doesn't exist
            FortranSyntaxError: If parsing fails
        """
        if not FPARSER2_AVAILABLE:
            raise ImportError("fparser2 is required for Fortran parsing")
        
        file_path = Path(file_path)
        if not file_path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                source = f.read()
            return self.build_from_source(source, str(file_path))
        except Exception as e:
            self.logger.error(f"Error reading file {file_path}: {e}")
            raise
    
    def build_from_source(self, source: str, filename: str = "<string>") -> FortranASTNode:
        """
        Build AST from Fortran source code.
        
        Args:
            source: Fortran source code
            filename: Source filename for error reporting
            
        Returns:
            Root AST node
        """
        if not FPARSER2_AVAILABLE:
            # Fallback to simple parsing without fparser2
            return self._simple_parse(source, filename)
        
        try:
            # Use fparser2 to parse Fortran
            reader = FortranStringReader(source)
            f2003_parser = Fortran2003.Program(reader)
            
            if f2003_parser is None:
                # Try as a module or procedure unit
                reader = FortranStringReader(source)
                f2003_parser = Fortran2003.Main_Program(reader)
                
                if f2003_parser is None:
                    reader = FortranStringReader(source)
                    f2003_parser = Fortran2003.Module(reader)
            
            if f2003_parser:
                return self._convert_fparser_node(f2003_parser, filename)
            else:
                self.logger.warning(f"Could not parse {filename} with fparser2, using fallback")
                return self._simple_parse(source, filename)
                
        except Exception as e:
            self.logger.warning(f"fparser2 parsing failed for {filename}: {e}, using fallback")
            return self._simple_parse(source, filename)
    
    def _convert_fparser_node(self, fparser_node, filename: str) -> FortranASTNode:
        """
        Convert fparser2 node to our AST representation.
        
        Args:
            fparser_node: fparser2 node
            filename: Source filename
            
        Returns:
            Converted AST node
        """
        node_type = type(fparser_node).__name__
        children = []
        attributes = {'filename': filename}
        
        # Extract source information if available
        if hasattr(fparser_node, 'item') and hasattr(fparser_node.item, 'span'):
            attributes['line_number'] = fparser_node.item.span[0]
        
        if hasattr(fparser_node, 'children'):
            for child in fparser_node.children:
                if child is not None:
                    children.append(self._convert_fparser_node(child, filename))
        elif hasattr(fparser_node, 'items'):
            for item in fparser_node.items:
                if item is not None:
                    children.append(self._convert_fparser_node(item, filename))
        
        # Extract specific attributes based on node type
        attributes.update(self._extract_node_attributes(fparser_node))
        
        return FortranASTNode(node_type, children, **attributes)
    
    def _extract_node_attributes(self, fparser_node) -> Dict[str, Any]:
        """Extract type-specific attributes from fparser2 node."""
        node_type = type(fparser_node).__name__
        attributes = {}
        
        # Handle different node types
        if node_type in ['Function_Subprogram', 'Subroutine_Subprogram']:
            if hasattr(fparser_node, 'children') and len(fparser_node.children) > 0:
                stmt = fparser_node.children[0]
                if hasattr(stmt, 'children') and len(stmt.children) > 1:
                    attributes['name'] = str(stmt.children[1])
        
        elif node_type == 'Program':
            if hasattr(fparser_node, 'children') and len(fparser_node.children) > 0:
                program_stmt = fparser_node.children[0]
                if hasattr(program_stmt, 'children') and len(program_stmt.children) > 1:
                    attributes['name'] = str(program_stmt.children[1])
        
        elif node_type == 'Module':
            if hasattr(fparser_node, 'children') and len(fparser_node.children) > 0:
                module_stmt = fparser_node.children[0]
                if hasattr(module_stmt, 'children') and len(module_stmt.children) > 1:
                    attributes['name'] = str(module_stmt.children[1])
        
        elif node_type == 'Call_Stmt':
            if hasattr(fparser_node, 'children') and len(fparser_node.children) > 0:
                attributes['target'] = str(fparser_node.children[0])
        
        elif node_type == 'Use_Stmt':
            if hasattr(fparser_node, 'children') and len(fparser_node.children) > 1:
                attributes['module'] = str(fparser_node.children[1])
        
        return attributes
    
    def _simple_parse(self, source: str, filename: str) -> FortranASTNode:
        """
        Simple fallback parser when fparser2 is not available.
        
        Args:
            source: Fortran source code
            filename: Source filename
            
        Returns:
            Basic AST representation
        """
        lines = source.split('\n')
        root = FortranASTNode('Program', filename=filename)
        
        current_unit = None
        line_number = 0
        
        for line_number, line in enumerate(lines, 1):
            line = line.strip()
            if not line or line.startswith('!') or line.startswith('C') or line.startswith('c'):
                continue
            
            # Simple pattern matching for basic constructs
            line_lower = line.lower()
            
            if line_lower.startswith('program '):
                name = line.split()[1] if len(line.split()) > 1 else 'unnamed'
                current_unit = FortranASTNode('Program', name=name, line_number=line_number)
                root.children.append(current_unit)
            
            elif line_lower.startswith('module '):
                name = line.split()[1] if len(line.split()) > 1 else 'unnamed'
                current_unit = FortranASTNode('Module', name=name, line_number=line_number)
                root.children.append(current_unit)
            
            elif line_lower.startswith('subroutine '):
                name = line.split()[1].split('(')[0] if len(line.split()) > 1 else 'unnamed'
                procedure = FortranASTNode('Subroutine', name=name, line_number=line_number)
                if current_unit:
                    current_unit.children.append(procedure)
                else:
                    root.children.append(procedure)
            
            elif line_lower.startswith('function '):
                name = line.split()[1].split('(')[0] if len(line.split()) > 1 else 'unnamed'
                procedure = FortranASTNode('Function', name=name, line_number=line_number)
                if current_unit:
                    current_unit.children.append(procedure)
                else:
                    root.children.append(procedure)
            
            elif line_lower.startswith('call '):
                target = line.split()[1].split('(')[0] if len(line.split()) > 1 else 'unknown'
                call_stmt = FortranASTNode('Call', target=target, line_number=line_number)
                if current_unit:
                    current_unit.children.append(call_stmt)
                else:
                    root.children.append(call_stmt)
            
            elif line_lower.startswith('use '):
                module_name = line.split()[1].split(',')[0] if len(line.split()) > 1 else 'unknown'
                use_stmt = FortranASTNode('Use', module=module_name, line_number=line_number)
                if current_unit:
                    current_unit.children.append(use_stmt)
                else:
                    root.children.append(use_stmt)
        
        return root
    
    def extract_procedures(self, ast_node: FortranASTNode) -> List[Dict[str, Any]]:
        """
        Extract procedure information from AST.
        
        Args:
            ast_node: Root AST node
            
        Returns:
            List of procedure dictionaries
        """
        procedures = []
        self._collect_procedures(ast_node, procedures)
        return procedures
    
    def _collect_procedures(self, node: FortranASTNode, procedures: List[Dict[str, Any]]):
        """Recursively collect procedure information."""
        if node.node_type in ['Function', 'Subroutine', 'Function_Subprogram', 'Subroutine_Subprogram']:
            proc_info = {
                'name': node.attributes.get('name', 'unnamed'),
                'type': 'function' if 'function' in node.node_type.lower() else 'subroutine',
                'line_number': node.line_number,
                'calls': []
            }
            
            # Extract calls within this procedure
            self._collect_calls(node, proc_info['calls'])
            procedures.append(proc_info)
        
        # Recurse into children
        for child in node.children:
            self._collect_procedures(child, procedures)
    
    def _collect_calls(self, node: FortranASTNode, calls: List[str]):
        """Recursively collect call statements."""
        if node.node_type in ['Call', 'Call_Stmt']:
            target = node.attributes.get('target', 'unknown')
            if target not in calls:
                calls.append(target)
        
        # Recurse into children
        for child in node.children:
            self._collect_calls(child, calls)
    
    def extract_modules(self, ast_node: FortranASTNode) -> List[Dict[str, Any]]:
        """
        Extract module information from AST.
        
        Args:
            ast_node: Root AST node
            
        Returns:
            List of module dictionaries
        """
        modules = []
        self._collect_modules(ast_node, modules)
        return modules
    
    def _collect_modules(self, node: FortranASTNode, modules: List[Dict[str, Any]]):
        """Recursively collect module information."""
        if node.node_type in ['Module']:
            module_info = {
                'name': node.attributes.get('name', 'unnamed'),
                'line_number': node.line_number,
                'procedures': [],
                'types': [],
                'uses': []
            }
            
            # Extract module contents
            for child in node.children:
                if child.node_type in ['Function', 'Subroutine']:
                    module_info['procedures'].append(child.attributes.get('name', 'unnamed'))
                elif child.node_type in ['Use', 'Use_Stmt']:
                    module_info['uses'].append(child.attributes.get('module', 'unknown'))
            
            modules.append(module_info)
        
        # Recurse into children
        for child in node.children:
            self._collect_modules(child, modules)